/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.metrics;



import com.bff.gaia.unified.model.pipeline.v1.MetricsApi.MonitoringInfo;

import com.bff.gaia.unified.runners.core.metrics.MetricsContainerImpl;

import com.bff.gaia.unified.runners.core.metrics.MetricsContainerStepMap;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.api.common.accumulators.Accumulator;

import com.bff.gaia.api.common.functions.RuntimeContext;

import com.bff.gaia.unified.sdk.metrics.DistributionResult;

import com.bff.gaia.unified.sdk.metrics.GaugeResult;

import com.bff.gaia.unified.sdk.metrics.MetricKey;

import com.bff.gaia.unified.sdk.metrics.MetricName;

import com.bff.gaia.unified.sdk.metrics.MetricQueryResults;

import com.bff.gaia.unified.sdk.metrics.MetricResult;

import com.bff.gaia.unified.sdk.metrics.MetricResults;

import com.bff.gaia.unified.sdk.metrics.MetricsFilter;

import com.bff.gaia.configuration.GlobalConfiguration;

import com.bff.gaia.configuration.MetricOptions;

import com.bff.gaia.metrics.Counter;

import com.bff.gaia.metrics.Gauge;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.util.HashMap;

import java.util.List;

import java.util.Map;



import static com.bff.gaia.unified.runners.core.metrics.MetricsContainerStepMap.asAttemptedOnlyMetricResults;



/**

 * Helper class for holding a {@link MetricsContainerImpl} and forwarding Unified metrics to Gaia

 * accumulators and metrics.

 */

public class GaiaMetricContainer {



  public static final String ACCUMULATOR_NAME = "__metricscontainers";



  private static final Logger LOG = LoggerFactory.getLogger(GaiaMetricContainer.class);



  private static final String METRIC_KEY_SEPARATOR =

      GlobalConfiguration.loadConfiguration().getString(MetricOptions.SCOPE_DELIMITER);



  private final RuntimeContext runtimeContext;

  private final Map<String, Counter> gaiaCounterCache;

  private final Map<String, GaiaDistributionGauge> gaiaDistributionGaugeCache;

  private final Map<String, GaiaGauge> gaiaGaugeCache;

  private final MetricsAccumulator metricsAccumulator;



  public GaiaMetricContainer(RuntimeContext runtimeContext) {

    this.runtimeContext = runtimeContext;

    this.gaiaCounterCache = new HashMap<>();

    this.gaiaDistributionGaugeCache = new HashMap<>();

    this.gaiaGaugeCache = new HashMap<>();



    Accumulator<MetricsContainerStepMap, MetricsContainerStepMap> metricsAccumulator =

        runtimeContext.getAccumulator(ACCUMULATOR_NAME);

    if (metricsAccumulator == null) {

      metricsAccumulator = new MetricsAccumulator();

      try {

        runtimeContext.addAccumulator(ACCUMULATOR_NAME, metricsAccumulator);

      } catch (Exception e) {

        LOG.error("Failed to create metrics accumulator.", e);

      }

    }

    this.metricsAccumulator = (MetricsAccumulator) metricsAccumulator;

  }



  public MetricsContainerImpl getMetricsContainer(String stepName) {

    return metricsAccumulator != null

        ? metricsAccumulator.getLocalValue().getContainer(stepName)

        : null;

  }



  /**

   * Update this container with metrics from the passed {@link MonitoringInfo}s, and send updates

   * along to Gaia's internal metrics framework.

   */

  public void updateMetrics(String stepName, List<MonitoringInfo> monitoringInfos) {

    getMetricsContainer(stepName).update(monitoringInfos);

    updateMetrics(stepName);

  }



  /**

   * Update Gaia's internal metrics ({@link this#gaiaCounterCache}) with the latest metrics for a

   * given step.

   */

  void updateMetrics(String stepName) {

    MetricResults metricResults = asAttemptedOnlyMetricResults(metricsAccumulator.getLocalValue());

    MetricQueryResults metricQueryResults =

        metricResults.queryMetrics(MetricsFilter.builder().addStep(stepName).build());

    updateCounters(metricQueryResults.getCounters());

    updateDistributions(metricQueryResults.getDistributions());

    updateGauge(metricQueryResults.getGauges());

  }



  private void updateCounters(Iterable<MetricResult<Long>> counters) {

    for (MetricResult<Long> metricResult : counters) {

      String gaiaMetricName = getGaiaMetricNameString(metricResult.getKey());



      Long update = metricResult.getAttempted();



      // update gaia metric

      Counter counter =

          gaiaCounterCache.computeIfAbsent(

              gaiaMetricName, n -> runtimeContext.getMetricGroup().counter(n));

      counter.dec(counter.getCount());

      counter.inc(update);

    }

  }



  private void updateDistributions(Iterable<MetricResult<DistributionResult>> distributions) {

    for (MetricResult<DistributionResult> metricResult : distributions) {

      String gaiaMetricName = getGaiaMetricNameString(metricResult.getKey());



      DistributionResult update = metricResult.getAttempted();



      // update gaia metric

      GaiaDistributionGauge gauge = gaiaDistributionGaugeCache.get(gaiaMetricName);

      if (gauge == null) {

        gauge =

            runtimeContext

                .getMetricGroup()

                .gauge(gaiaMetricName, new GaiaDistributionGauge(update));

        gaiaDistributionGaugeCache.put(gaiaMetricName, gauge);

      } else {

        gauge.update(update);

      }

    }

  }



  private void updateGauge(Iterable<MetricResult<GaugeResult>> gauges) {

    for (MetricResult<GaugeResult> metricResult : gauges) {

      String gaiaMetricName = getGaiaMetricNameString(metricResult.getKey());



      GaugeResult update = metricResult.getAttempted();



      // update gaia metric

      GaiaGauge gauge = gaiaGaugeCache.get(gaiaMetricName);

      if (gauge == null) {

        gauge = runtimeContext.getMetricGroup().gauge(gaiaMetricName, new GaiaGauge(update));

        gaiaGaugeCache.put(gaiaMetricName, gauge);

      } else {

        gauge.update(update);

      }

    }

  }



  @VisibleForTesting

  static String getGaiaMetricNameString(MetricKey metricKey) {

    MetricName metricName = metricKey.metricName();

    // We use only the MetricName here, the step name is already contained

    // in the operator name which is passed to Gaia's MetricGroup to which

    // the metric with the following name will be added.

    return metricName.getNamespace() + METRIC_KEY_SEPARATOR + metricName.getName();

  }



  /** Gaia {@link Gauge} for {@link DistributionResult}. */

  public static class GaiaDistributionGauge implements Gauge<DistributionResult> {



    DistributionResult data;



    GaiaDistributionGauge(DistributionResult data) {

      this.data = data;

    }



    void update(DistributionResult data) {

      this.data = data;

    }



    @Override

    public DistributionResult getValue() {

      return data;

    }

  }



  /** Gaia {@link Gauge} for {@link GaugeResult}. */

  public static class GaiaGauge implements Gauge<GaugeResult> {



    GaugeResult data;



    GaiaGauge(GaugeResult data) {

      this.data = data;

    }



    void update(GaugeResult update) {

      this.data = update;

    }



    @Override

    public GaugeResult getValue() {

      return data;

    }

  }

}